from state_machine import StateMachine
import sensor, image, time, os, tf


net = "trained.tflite"
labels = [line.rstrip('\n') for line in open("labels.txt")]

clock = time.clock()

# 有限状态集合
single = ["Start", "Patrol", "Pause", "Recognition", "Return"]
multiple = []

target = None

# 自定义状态转变函数
def start_transitions(txt):
    print("start")
    global target
    splitted_txt = txt.split(None, 1)  
    word, txt = splitted_txt if len(splitted_txt) > 1 else (txt,"")
    if word == "Single":   
        newState = "single_state" # 如果第一个词是Single则可转换到"Single状态"
    else:
        newState = "error_state"  # 如果第一个词不是Single则进入终止状态

    sensor.reset()
    sensor.set_pixformat(sensor.GRAYSCALE) # 灰度更快(160x120 max on OpenMV-M7)
    sensor.set_framesize(sensor.QQVGA)
    sensor.skip_frames(time = 2000)
    while(True):
        clock.tick()
        img = sensor.snapshot()
        img.lens_corr()

        # 下面的`threshold`应设置为足够高的值，以滤除在图像中检测到的具有
        # 低边缘幅度的噪声矩形。最适用与背景形成鲜明对比的矩形。
        for r in img.find_rects(threshold = 10000):
            img.draw_rectangle(r.rect(), color = (255, 0, 0))
            # 数字识别
            for obj in tf.classify(net, img.crop([r[0],r[1],r[2],r[3]]), min_scale=1.0, scale_mul=0.8, x_overlap=0.5, y_overlap=0.5):
                print("**********\nPredictions at [x=%d,y=%d,w=%d,h=%d]" % obj.rect())
                img.draw_rectangle(obj.rect())
                # This combines the labels and confidence values into a list of tuples
                predictions_list = list(zip(labels, obj.output()))

                max_sml = 0
                for i in range(len(predictions_list)):
                    if predictions_list[i][1] > max_sml:
                        max_sml = predictions_list[i][1]
                        target = predictions_list[i][0]
                    print("%s = %f" % (predictions_list[i][0], predictions_list[i][1]))
                break
            return (newState, txt)

    return (newState, txt)        # 返回新状态和余下的语句txt

def single_state_transitions(txt):
    print("single")
    splitted_txt = txt.split(None,1)
    word, txt = splitted_txt if len(splitted_txt) > 1 else (txt,"")
    if word == "Patrol":
        newState = "patrol_state"
    else:
        newState = "error_state"
    return (newState, txt)

def patrol_state_transitions(txt):
    print("patrol")
    splitted_txt = txt.split(None,1)
    word, txt = splitted_txt if len(splitted_txt) > 1 else (txt,"")
    if word == "Pause":
        newState = "pause_state"
    else:
        newState = "error_state"
    return (newState, txt)

def pause_state_transitions(txt):
    print("pause")
    splitted_txt = txt.split(None,1)
    word, txt = splitted_txt if len(splitted_txt) > 1 else (txt,"")
    if word == "Recognition":
        newState = "recognition_state"
    else:
        newState = "error_state"
    return (newState, txt)

def recognition_state_transitions(txt):
    print("recognition")
    global target
    splitted_txt = txt.split(None,1)
    word, txt = splitted_txt if len(splitted_txt) > 1 else (txt,"")
    if word in single:
        newState = "end_state"
    else:
        newState = "error_state"

    sensor.set_pixformat(sensor.GRAYSCALE)
    while(True):
        clock.tick()
        img = sensor.snapshot()
        # 镜头畸变校正
        img.lens_corr()

        # 下面的`threshold`应设置为足够高的值，以滤除在图像中检测到的具有
        # 低边缘幅度的噪声矩形。最适用与背景形成鲜明对比的矩形。
        rects = img.find_rects(threshold = 10000)
        rects = sorted(rects, key=(lambda r:r.x()))
        print(rects)
        for r in rects:
            #img.draw_rectangle(r.rect(), color = (255, 0, 0))
            # 数字识别
            for obj in tf.classify(net, img,[r[0] - 8,r[1] - 8,r[2] + 16,r[3] + 16], min_scale=1.0, scale_mul=0.8, x_overlap=0.5, y_overlap=0.5):
                #print("**********\nPredictions at [x=%d,y=%d,w=%d,h=%d]" % obj.rect())
                img.draw_rectangle(obj.rect())
                # This combines the labels and confidence values into a list of tuples
                predictions_list = list(zip(labels, obj.output()))

                max_sml = 0
                max_pre_dic = None
                for i in range(len(predictions_list)):
                    if predictions_list[i][1] > max_sml:
                        max_sml = predictions_list[i][1]
                        max_pre_dic = predictions_list[i][0]
                print(str(max_pre_dic) + ':' + str(max_sml) + ';')
                # for i in range(len(predictions_list)):
                    # if (predictions_list[i][1] > 0.75) and (predictions_list[i][0] == target):
                    #     print("detected")
                    #     return (newState, txt)
                    #print("%s = %f" % (predictions_list[i][0], predictions_list[i][1]))
                #break
    return (newState, txt)

if __name__== "__main__":
    m = StateMachine()
    m.add_state("Start", start_transitions)      # 添加初始状态
    m.add_state("Single_state", single_state_transitions)
    m.add_state("Patrol_state", patrol_state_transitions)
    m.add_state("Pause_state", pause_state_transitions)
    m.add_state("Recognition_state", recognition_state_transitions)
    m.add_state("end_state", None, end_state=1)  # 添加最终状态
    m.add_state("error_state", None, end_state=1)
    
    m.set_start("Start") # 设置开始状态
    m.run("Single Patrol Pause Recognition")
    print("end")
